import requests
from PIL import Image
from io import BytesIO
import base64
import uuid
import time
import zipfile
from fastapi.responses import Response, JSONResponse

class JiugongAPI:
    def __init__(self):
        pass
        
    def handle(self, url: str, type: str = "json"):
        try:
            # 下载远程图片
            response = requests.get(url)
            response.raise_for_status()  # 确保请求成功
            
            # 使用PIL打开图片
            img = Image.open(BytesIO(response.content))
            
            # 获取图片尺寸
            width, height = img.size
            
            # 计算每个小格子的尺寸
            grid_width = width // 3
            grid_height = height // 3
            
            # 创建9个子图片
            grid_images = []
            for i in range(3):
                for j in range(3):
                    # 计算裁剪区域
                    left = j * grid_width
                    upper = i * grid_height
                    right = left + grid_width
                    lower = upper + grid_height
                    
                    # 裁剪图片
                    grid_img = img.crop((left, upper, right, lower))
                    grid_images.append(grid_img)
            
            # 根据type参数决定返回类型
            if type.lower() == "json":
                # 将处理后的图片转为base64编码
                base64_images = []
                for idx, grid_img in enumerate(grid_images):
                    buffered = BytesIO()
                    grid_img.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()
                    base64_images.append({
                        "index": idx + 1,
                        "image_base64": f"data:image/png;base64,{img_str}"
                    })
                
                return JSONResponse(content={
                    "success": True,
                    "code": 200,
                    "message": "图片已成功分割为9宫格",
                    "data": {
                        "images": base64_images
                    }
                }, status_code=200, media_type="application/json, charset=utf-8")
            elif type.lower() == "image":
                # 创建内存中的ZIP文件
                zip_buffer = BytesIO()
                
                with zipfile.ZipFile(zip_buffer, 'w') as zipf:
                    # 将每个图片添加到ZIP文件
                    for idx, grid_img in enumerate(grid_images):
                        # 将图片转换为字节流
                        img_buffer = BytesIO()
                        grid_img.save(img_buffer, format="PNG")
                        img_buffer.seek(0)
                        
                        # 添加到ZIP文件
                        zipf.writestr(f"grid_{idx+1}.png", img_buffer.getvalue())
                
                # 准备返回ZIP文件
                zip_buffer.seek(0)
                zip_filename = f"jiugong_{uuid.uuid4().hex}_{int(time.time())}.zip"
                
                # 直接返回内存中的ZIP文件
                return Response(
                    content=zip_buffer.getvalue(),
                    media_type="application/zip",
                    status_code=200,
                    headers={
                        "Content-Disposition": f"attachment; filename={zip_filename}"
                    }
                )
            else:
                return {
                    "success": False,
                    "message": f"不支持的type参数: {type}，支持的类型为json或image",
                    "data": None
                }
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "message": f"获取远程图片失败: {str(e)}",
                "data": None
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"处理图片时出错: {str(e)}",
                "data": None
            }